import numpy as np
from scipy import stats


def calculate_power_greenhouse_geisser(alpha, M, V, inerr, sphericity, g1, power_target):
    """
    Calculate sample size needed for desired power in repeated measures ANOVA with Greenhouse-Geisser correction

    Parameters:
    -----------
    alpha : float
        Significance level (0-1)
    M : int
        Number of levels (groups)
    V : float
        Variance of means
    inerr : float
        Within-group error term
    sphericity : float
        Measure of sphericity (epsilon) (0-1)
    g1 : float
        Bias term multiplier
    power_target : float
        Desired power (1-100)

    Returns:
    --------
    dict
        Dictionary containing calculation results
    """
    # Input validation
    if not 0 < alpha < 1:
        raise ValueError("Test significance level must be in 0-1")
    if not 0 <= sphericity <= 1:
        raise ValueError("Measure of 'sphericity' must be in 0-1")
    if M <= 0 or not isinstance(M, int):
        raise ValueError("The Number of levels must be positive integer")
    if not 1 <= power_target <= 100:
        raise ValueError("Power % must be in 1-100")

    # Calculate effect size
    effect_size = V / (inerr ** 2)

    # Initial sample size
    n = M + 1
    power_achieved = 0

    # Iterate until desired power is achieved
    while power_achieved < power_target / 100:
        # Calculate degrees of freedom
        df1 = (M - 1) * (sphericity + g1 / (n - 1))
        df2 = (n - 1) * (M - 1) * (sphericity + g1 / (n - 1))

        # Calculate non-centrality parameter
        ncp = n * M * effect_size * (sphericity + g1 / (n - 1))

        # Calculate critical F value
        f_crit = stats.f.ppf(1 - alpha, df1, df2)

        # Calculate achieved power
        power_achieved = 1 - stats.ncf.cdf(f_crit, df1, df2, ncp)

        n += 0.01

    # Adjust final sample size
    n = int(n - 0.01)

    return {
        'test_significance_level': alpha,
        'variance_of_means': V,
        'within_group_error': inerr,
        'sphericity': sphericity,
        'bias_term_multiplier': g1,
        'effect_size': effect_size,
        'power_percent': power_target,
        'sample_size': n
    }


# Test with example from the paper
def test_example():
    result = calculate_power_greenhouse_geisser(
        alpha=0.05,
        M=4,
        V=2.813,
        inerr=6.11,
        sphericity=0.88,
        g1=-1.98,
        power_target=90
    )

    print("Results:")
    print(f"Test Significance level: {result['test_significance_level']}")
    print(f"Variance of means: {result['variance_of_means']}")
    print(f"Within-group error term: {result['within_group_error']}")
    print(f"Measure of sphericity: {result['sphericity']}")
    print(f"Bias term multiplier: {result['bias_term_multiplier']}")
    print(f"Effect size: {result['effect_size']:.4f}")
    print(f"Power %: {result['power_percent']}")
    print(f"Required sample size (n): {result['sample_size']}")


if __name__ == "__main__":
    test_example()